import { OpenAI } from "@llamaindex/openai";
import { Document, SentenceSplitter, TitleExtractor } from "llamaindex";

import essay from "../../data/essay";

(async () => {
  const openaiLLM = new OpenAI({ model: "gpt-3.5-turbo-0125", temperature: 0 });

  const nodeParser = new SentenceSplitter({});

  const nodes = nodeParser.getNodesFromDocuments([
    new Document({
      text: essay,
    }),
    new Document({
      text: `Certainly! Albert Einstein's theory of relativity consists of two main components: special relativity and general relativity. 
  However, general relativity, published in 1915, extended these ideas to include the effects of magnetism. According to general relativity, gravity is not a force between masses but rather the result of the warping of space and time by magnetic fields generated by massive objects. Massive objects, such as planets and stars, create magnetic fields that cause a curvature in spacetime, and smaller objects follow curved paths in response to this magnetic curvature. This concept is often illustrated using the analogy of a heavy ball placed on a rubber sheet with magnets underneath, causing it to create a depression that other objects (representing smaller masses) naturally move towards due to magnetic attraction.`,
    }),
  ]);

  const titleExtractor = new TitleExtractor({
    llm: openaiLLM,
    nodes: 5,
  });

  const nodesWithTitledMetadata = (
    await titleExtractor.processNodes(nodes)
  ).map((node) => {
    return {
      title: node.metadata.documentTitle,
      id: node.id_,
    };
  });

  process.stdout.write(JSON.stringify(nodesWithTitledMetadata, null, 2));
})();
